"""
任务存储服务

提供任务的存储、查询和管理功能
使用内存存储，支持并发访问
"""

import asyncio
from datetime import datetime
from typing import Optional, Dict
from loguru import logger

from src.schemas.task import Task, TaskType, TaskStatus


class TaskStore:
    """任务存储"""
    
    _instance = None
    _lock = asyncio.Lock()
    
    def __new__(cls):
        """单例模式"""
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        """初始化任务存储"""
        if self._initialized:
            return
        
        self._tasks: Dict[str, Task] = {}
        self._task_lock = asyncio.Lock()
        self._initialized = True
        
        logger.info("任务存储初始化完成")
    
    async def create_task(
        self,
        task_id: str,
        task_type: TaskType,
        context: dict = None,
        files_hash: str = None
    ) -> Task:
        """
        创建新任务
        
        Args:
            task_id: 任务ID
            task_type: 任务类型
            context: 任务上下文数据
            files_hash: 附件哈希值（用于缓存查询）
            
        Returns:
            Task: 创建的任务
        """
        async with self._task_lock:
            if task_id in self._tasks:
                existing_task = self._tasks[task_id]
                logger.warning(f"任务已存在: {task_id}")
                
                # 更新 files_hash（如果提供了新的且与现有不同）
                if files_hash and existing_task.files_hash != files_hash:
                    old_hash = existing_task.files_hash[:16] + "..." if existing_task.files_hash else "None"
                    new_hash = files_hash[:16] + "..."
                    logger.info(f"更新任务的 files_hash: {task_id}, 旧={old_hash}, 新={new_hash}")
                    existing_task.files_hash = files_hash
                    existing_task.updated_at = datetime.now()
                
                return existing_task
            
            now = datetime.now()
            task = Task(
                task_id=task_id,
                task_type=task_type,
                status=TaskStatus.PENDING,
                progress=0.0,
                created_at=now,
                updated_at=now,
                context=context or {},
                files_hash=files_hash
            )
            
            self._tasks[task_id] = task
            
            if files_hash:
                logger.info(f"创建任务: {task_id}, 类型: {task_type}, 文件哈希: {files_hash[:16]}...")
            else:
                logger.info(f"创建任务: {task_id}, 类型: {task_type}")
            
            return task
    
    async def get_task(self, task_id: str) -> Optional[Task]:
        """
        获取任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            Optional[Task]: 任务对象，不存在返回None
        """
        async with self._task_lock:
            return self._tasks.get(task_id)
    
    async def update_task(
        self,
        task_id: str,
        status: Optional[TaskStatus] = None,
        progress: Optional[float] = None,
        result: Optional[any] = None,
        error_message: Optional[str] = None,
        context: Optional[dict] = None
    ) -> Optional[Task]:
        """
        更新任务状态
        
        Args:
            task_id: 任务ID
            status: 新状态
            progress: 新进度
            result: 任务结果
            error_message: 错误信息
            context: 任务上下文数据（会合并到现有context中）
            
        Returns:
            Optional[Task]: 更新后的任务，不存在返回None
        """
        async with self._task_lock:
            task = self._tasks.get(task_id)
            if not task:
                logger.warning(f"任务不存在: {task_id}")
                return None
            
            # 如果任务已取消，不允许更新
            if task.cancelled and status != TaskStatus.CANCELLED:
                logger.warning(f"任务已取消，不允许更新: {task_id}")
                return task
            
            now = datetime.now()
            
            if status is not None:
                task.status = status
            
            if progress is not None:
                task.progress = progress
            
            if result is not None:
                task.result = result
            
            if error_message is not None:
                task.error_message = error_message
            
            if context is not None:
                # 合并context（不覆盖，而是更新）
                task.context.update(context)
            
            task.updated_at = now
            
            # 如果任务完成或失败，设置完成时间
            if status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.TIMEOUT, TaskStatus.CANCELLED]:
                if task.completed_at is None:
                    task.completed_at = now
            
            logger.debug(f"更新任务: {task_id}, 状态: {status}, 进度: {progress}")
            
            return task
    
    async def cancel_task(self, task_id: str) -> bool:
        """
        取消任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            bool: 是否成功取消
        """
        async with self._task_lock:
            task = self._tasks.get(task_id)
            if not task:
                logger.warning(f"任务不存在: {task_id}")
                return False
            
            # 只有进行中或等待中的任务可以取消
            if task.status in [TaskStatus.PENDING, TaskStatus.PROCESSING]:
                task.cancelled = True
                task.status = TaskStatus.CANCELLED
                task.updated_at = datetime.now()
                task.completed_at = datetime.now()
                
                logger.info(f"取消任务: {task_id}")
                return True
            else:
                logger.warning(f"任务状态为 {task.status}，无法取消: {task_id}")
                return False
    
    async def delete_task(self, task_id: str) -> bool:
        """
        删除任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            bool: 是否成功删除
        """
        async with self._task_lock:
            if task_id in self._tasks:
                del self._tasks[task_id]
                logger.info(f"删除任务: {task_id}")
                return True
            return False
    
    async def list_tasks(
        self,
        task_type: Optional[TaskType] = None,
        status: Optional[TaskStatus] = None,
        limit: int = 100
    ) -> list[Task]:
        """
        列出任务
        
        Args:
            task_type: 任务类型过滤
            status: 状态过滤
            limit: 最大返回数量
            
        Returns:
            list[Task]: 任务列表
        """
        async with self._task_lock:
            tasks = list(self._tasks.values())
            
            # 过滤
            if task_type is not None:
                tasks = [t for t in tasks if t.task_type == task_type]
            
            if status is not None:
                tasks = [t for t in tasks if t.status == status]
            
            # 按创建时间倒序排序
            tasks.sort(key=lambda t: t.created_at, reverse=True)
            
            return tasks[:limit]
    
    async def find_task_by_files_hash(
        self,
        files_hash: str,
        task_type: Optional[TaskType] = None,
        only_completed: bool = True
    ) -> Optional[Task]:
        """
        根据附件哈希查找历史任务
        
        用于实现结果缓存：如果之前已经处理过相同的附件集合，
        可以直接返回历史结果，避免重复计算
        
        Args:
            files_hash: 附件哈希值
            task_type: 任务类型过滤（可选）
            only_completed: 是否只返回已完成的任务
            
        Returns:
            Optional[Task]: 找到的任务，按创建时间倒序返回最新的一个，未找到返回None
        """
        async with self._task_lock:
            if not files_hash:
                return None
            
            # 查找匹配的任务
            matching_tasks = []
            
            for task in self._tasks.values():
                # 检查files_hash是否匹配
                if task.files_hash != files_hash:
                    continue
                
                # 检查任务类型
                if task_type is not None and task.task_type != task_type:
                    continue
                
                # 检查是否只返回已完成的任务
                if only_completed and task.status != TaskStatus.COMPLETED:
                    continue
                
                matching_tasks.append(task)
            
            if not matching_tasks:
                logger.debug(f"未找到文件哈希为 {files_hash[:16]}... 的历史任务")
                return None
            
            # 按创建时间倒序排序，返回最新的
            matching_tasks.sort(key=lambda t: t.created_at, reverse=True)
            latest_task = matching_tasks[0]
            
            logger.info(
                f"找到文件哈希 {files_hash[:16]}... 的历史任务: {latest_task.task_id}, "
                f"状态: {latest_task.status}, 创建时间: {latest_task.created_at}"
            )
            
            return latest_task
    
    async def cleanup_old_tasks(self, max_age_seconds: int = 3600):
        """
        清理旧任务
        
        Args:
            max_age_seconds: 最大保留时间（秒）
        """
        async with self._task_lock:
            now = datetime.now()
            to_delete = []
            
            for task_id, task in self._tasks.items():
                age = (now - task.created_at).total_seconds()
                if age > max_age_seconds and task.status in [
                    TaskStatus.COMPLETED,
                    TaskStatus.FAILED,
                    TaskStatus.TIMEOUT,
                    TaskStatus.CANCELLED
                ]:
                    to_delete.append(task_id)
            
            for task_id in to_delete:
                del self._tasks[task_id]
            
            if to_delete:
                logger.info(f"清理 {len(to_delete)} 个旧任务")


# 全局单例
task_store = TaskStore()


def get_task_store() -> TaskStore:
    """获取任务存储实例"""
    return task_store

